{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "a04259d2",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#RRDBNet_arch.py\n",
    "\n",
    "import functools\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "\n",
    "def make_layer(block, n_layers):\n",
    "    layers = []\n",
    "    for _ in range(n_layers):\n",
    "        layers.append(block())\n",
    "    return nn.Sequential(*layers)\n",
    "\n",
    "\n",
    "class ResidualDenseBlock_5C(nn.Module):\n",
    "    def __init__(self, nf=64, gc=32, bias=True):\n",
    "        super(ResidualDenseBlock_5C, self).__init__()\n",
    "        # gc: growth channel, i.e. intermediate channels\n",
    "        self.conv1 = nn.Conv2d(nf, gc, 3, 1, 1, bias=bias)\n",
    "        self.conv2 = nn.Conv2d(nf + gc, gc, 3, 1, 1, bias=bias)\n",
    "        self.conv3 = nn.Conv2d(nf + 2 * gc, gc, 3, 1, 1, bias=bias)\n",
    "        self.conv4 = nn.Conv2d(nf + 3 * gc, gc, 3, 1, 1, bias=bias)\n",
    "        self.conv5 = nn.Conv2d(nf + 4 * gc, nf, 3, 1, 1, bias=bias)\n",
    "        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)\n",
    "\n",
    "        # initialization\n",
    "        # mutil.initialize_weights([self.conv1, self.conv2, self.conv3, self.conv4, self.conv5], 0.1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x1 = self.lrelu(self.conv1(x))\n",
    "        x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1)))\n",
    "        x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1)))\n",
    "        x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1)))\n",
    "        x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))\n",
    "        return x5 * 0.2 + x\n",
    "\n",
    "\n",
    "class RRDB(nn.Module):\n",
    "    '''Residual in Residual Dense Block'''\n",
    "\n",
    "    def __init__(self, nf, gc=32):\n",
    "        super(RRDB, self).__init__()\n",
    "        self.RDB1 = ResidualDenseBlock_5C(nf, gc)\n",
    "        self.RDB2 = ResidualDenseBlock_5C(nf, gc)\n",
    "        self.RDB3 = ResidualDenseBlock_5C(nf, gc)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.RDB1(x)\n",
    "        out = self.RDB2(out)\n",
    "        out = self.RDB3(out)\n",
    "        return out * 0.2 + x\n",
    "\n",
    "\n",
    "class RRDBNet(nn.Module):\n",
    "    def __init__(self, in_nc, out_nc, nf, nb, gc=32):\n",
    "        super(RRDBNet, self).__init__()\n",
    "        RRDB_block_f = functools.partial(RRDB, nf=nf, gc=gc)\n",
    "\n",
    "        self.conv_first = nn.Conv2d(in_nc, nf, 3, 1, 1, bias=True)\n",
    "        self.RRDB_trunk = make_layer(RRDB_block_f, nb)\n",
    "        self.trunk_conv = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)\n",
    "        #### upsampling\n",
    "        self.upconv1 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)\n",
    "        self.upconv2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)\n",
    "        self.HRconv = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)\n",
    "        self.conv_last = nn.Conv2d(nf, out_nc, 3, 1, 1, bias=True)\n",
    "\n",
    "        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)\n",
    "\n",
    "    def forward(self, x):\n",
    "        fea = self.conv_first(x)\n",
    "        trunk = self.trunk_conv(self.RRDB_trunk(fea))\n",
    "        fea = fea + trunk\n",
    "\n",
    "        fea = self.lrelu(self.upconv1(F.interpolate(fea, scale_factor=2, mode='nearest')))\n",
    "        fea = self.lrelu(self.upconv2(F.interpolate(fea, scale_factor=2, mode='nearest')))\n",
    "        out = self.conv_last(self.lrelu(self.HRconv(fea)))\n",
    "\n",
    "        return out\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "70002df6",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "###################################\n",
      "\n",
      "Saving to  ./models/RRDB_ESRGAN_x4.pth\n"
     ]
    }
   ],
   "source": [
    "#transer_RRDB_models.py\n",
    "\n",
    "import os\n",
    "import torch\n",
    "import RRDBNet_arch as arch\n",
    "\n",
    "pretrained_net = torch.load('./models/RRDB_ESRGAN_x4_old_arch.pth')\n",
    "save_path = './models/RRDB_ESRGAN_x4.pth'\n",
    "\n",
    "crt_model = arch.RRDBNet(3, 3, 64, 23, gc=32)\n",
    "crt_net = crt_model.state_dict()\n",
    "\n",
    "load_net_clean = {}\n",
    "for k, v in pretrained_net.items():\n",
    "    if k.startswith('module.'):\n",
    "        load_net_clean[k[7:]] = v\n",
    "    else:\n",
    "        load_net_clean[k] = v\n",
    "pretrained_net = load_net_clean\n",
    "\n",
    "print('###################################\\n')\n",
    "tbd = []\n",
    "for k, v in crt_net.items():\n",
    "    tbd.append(k)\n",
    "\n",
    "# directly copy\n",
    "for k, v in crt_net.items():\n",
    "    if k in pretrained_net and pretrained_net[k].size() == v.size():\n",
    "        crt_net[k] = pretrained_net[k]\n",
    "        tbd.remove(k)\n",
    "\n",
    "crt_net['conv_first.weight'] = pretrained_net['model.0.weight']\n",
    "crt_net['conv_first.bias'] = pretrained_net['model.0.bias']\n",
    "\n",
    "for k in tbd.copy():\n",
    "    if 'RDB' in k:\n",
    "        ori_k = k.replace('RRDB_trunk.', 'model.1.sub.')\n",
    "        if '.weight' in k:\n",
    "            ori_k = ori_k.replace('.weight', '.0.weight')\n",
    "        elif '.bias' in k:\n",
    "            ori_k = ori_k.replace('.bias', '.0.bias')\n",
    "        crt_net[k] = pretrained_net[ori_k]\n",
    "        tbd.remove(k)\n",
    "\n",
    "crt_net['trunk_conv.weight'] = pretrained_net['model.1.sub.23.weight']\n",
    "crt_net['trunk_conv.bias'] = pretrained_net['model.1.sub.23.bias']\n",
    "crt_net['upconv1.weight'] = pretrained_net['model.3.weight']\n",
    "crt_net['upconv1.bias'] = pretrained_net['model.3.bias']\n",
    "crt_net['upconv2.weight'] = pretrained_net['model.6.weight']\n",
    "crt_net['upconv2.bias'] = pretrained_net['model.6.bias']\n",
    "crt_net['HRconv.weight'] = pretrained_net['model.8.weight']\n",
    "crt_net['HRconv.bias'] = pretrained_net['model.8.bias']\n",
    "crt_net['conv_last.weight'] = pretrained_net['model.10.weight']\n",
    "crt_net['conv_last.bias'] = pretrained_net['model.10.bias']\n",
    "\n",
    "torch.save(crt_net, save_path)\n",
    "print('Saving to ', save_path)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "89668f3a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model path models/RRDB_PSNR_x4.pth. \n",
      "Testing...\n",
      "1 00003\n",
      "2 00017_gray\n",
      "3 0014\n",
      "4 0030\n",
      "5 ADE_val_00000114\n",
      "6 baboon\n",
      "7 children-alpha\n",
      "8 comic\n",
      "9 OST_009\n",
      "10 tree_alpha_16bit\n",
      "11 wolf_gray\n"
     ]
    }
   ],
   "source": [
    "#test.py\n",
    "\n",
    "import os.path as osp\n",
    "import glob\n",
    "import cv2\n",
    "import numpy as np\n",
    "import torch\n",
    "import RRDBNet_arch as arch\n",
    "\n",
    "model_path = 'models/RRDB_PSNR_x4.pth'  # models/RRDB_ESRGAN_x4.pth OR models/RRDB_PSNR_x4.pth\n",
    "#device = torch.device('cuda')  # if you want to run on CPU, change 'cuda' -> cpu\n",
    "device = torch.device('cpu')\n",
    "\n",
    "test_img_folder = 'LR/*'\n",
    "\n",
    "model = arch.RRDBNet(3, 3, 64, 23, gc=32)\n",
    "model.load_state_dict(torch.load(model_path), strict=True)\n",
    "model.eval()\n",
    "model = model.to(device)\n",
    "\n",
    "print('Model path {:s}. \\nTesting...'.format(model_path))\n",
    "\n",
    "idx = 0\n",
    "for path in glob.glob(test_img_folder):\n",
    "    idx += 1\n",
    "    base = osp.splitext(osp.basename(path))[0]\n",
    "    print(idx, base)\n",
    "    # read images\n",
    "    img = cv2.imread(path, cv2.IMREAD_COLOR)\n",
    "    img = img * 1.0 / 255\n",
    "    img = torch.from_numpy(np.transpose(img[:, :, [2, 1, 0]], (2, 0, 1))).float()\n",
    "    img_LR = img.unsqueeze(0)\n",
    "    img_LR = img_LR.to(device)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        output = model(img_LR).data.squeeze().float().cpu().clamp_(0, 1).numpy()\n",
    "    output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0))\n",
    "    output = (output * 255.0).round()\n",
    "    cv2.imwrite('results/{:s}_PSNR.png'.format(base), output)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "b4b4f722",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model path models/RRDB_ESRGAN_x4.pth. \n",
      "Testing...\n",
      "1 00003\n",
      "2 00017_gray\n",
      "3 0014\n",
      "4 0030\n",
      "5 ADE_val_00000114\n",
      "6 baboon\n",
      "7 children-alpha\n",
      "8 comic\n",
      "9 OST_009\n",
      "10 tree_alpha_16bit\n",
      "11 wolf_gray\n"
     ]
    }
   ],
   "source": [
    "import os.path as osp\n",
    "import glob\n",
    "import cv2\n",
    "import numpy as np\n",
    "import torch\n",
    "import RRDBNet_arch as arch\n",
    "\n",
    "model_path = 'models/RRDB_ESRGAN_x4.pth'  # models/RRDB_ESRGAN_x4.pth OR models/RRDB_PSNR_x4.pth\n",
    "#device = torch.device('cuda')  # if you want to run on CPU, change 'cuda' -> cpu\n",
    "device = torch.device('cpu')\n",
    "\n",
    "test_img_folder = 'LR/*'\n",
    "\n",
    "model = arch.RRDBNet(3, 3, 64, 23, gc=32)\n",
    "model.load_state_dict(torch.load(model_path), strict=True)\n",
    "model.eval()\n",
    "model = model.to(device)\n",
    "\n",
    "print('Model path {:s}. \\nTesting...'.format(model_path))\n",
    "\n",
    "idx = 0\n",
    "for path in glob.glob(test_img_folder):\n",
    "    idx += 1\n",
    "    base = osp.splitext(osp.basename(path))[0]\n",
    "    print(idx, base)\n",
    "    # read images\n",
    "    img = cv2.imread(path, cv2.IMREAD_COLOR)\n",
    "    img = img * 1.0 / 255\n",
    "    img = torch.from_numpy(np.transpose(img[:, :, [2, 1, 0]], (2, 0, 1))).float()\n",
    "    img_LR = img.unsqueeze(0)\n",
    "    img_LR = img_LR.to(device)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        output = model(img_LR).data.squeeze().float().cpu().clamp_(0, 1).numpy()\n",
    "    output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0))\n",
    "    output = (output * 255.0).round()\n",
    "    cv2.imwrite('results/{:s}_ESRGAN.png'.format(base), output)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
